# ============================================
# 多代理协作示例代码整理版
# ============================================

# -----------------------------
# 基础模块导入
# -----------------------------
import functools
import operator
import os
from typing import Annotated, TypedDict, Sequence, Literal

from dotenv import load_dotenv  # 用于加载 .env 文件中的 API key 等配置

# -----------------------------
# LangChain / LangGraph 导入
# -----------------------------
# 消息类型
from langchain_core.messages import BaseMessage, AIMessage, ToolMessage, HumanMessage
# 聊天提示模板
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
# 工具装饰器
from langchain_core.tools import tool
# Python 执行工具
from langchain_experimental.utilities import PythonREPL
# Tavily 搜索工具
from langchain_community.tools.tavily_search import TavilySearchResults
# OpenAI LLM 接口
from langchain_openai import ChatOpenAI

# 状态图相关
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from langgraph.prebuilt import ToolNode

# -----------------------------
# 环境变量加载
# -----------------------------
load_dotenv()  # 从 .env 文件中读取 API Key

# ============================================
# 1. 代理创建函数
# ============================================

def create_agent(llm, tools, system_message: str):
    """
    创建一个代理，绑定 LLM 和工具，并生成带有系统消息的提示模板。
    系统消息会告诉代理如何协作和使用工具。
    """
    # 创建聊天提示模板
    prompt = ChatPromptTemplate.from_messages(
        [
            (
                "system",
                "你是一个有帮助的AI助手，与其他助手合作。"
                " 使用提供的工具推进问题回答。"
                " 如果你不能完全回答，另一个拥有不同工具的助手会接着你的位置继续帮助。"
                " 如果你或其他助手有最终答案，请在回答前加上FINAL ANSWER。"
                " 你可以使用以下工具: {tool_names}。\n{system_message}",
            ),
            # 占位符，用于存储对话历史
            MessagesPlaceholder(variable_name="messages"),
        ]
    )

    # 使用 partial 填充系统消息和工具名称
    prompt = prompt.partial(system_message=system_message)
    prompt = prompt.partial(tool_names=", ".join([tool.name for tool in tools]))

    # 绑定工具并返回 prompt -> llm 流水线
    return prompt | llm.bind_tools(tools)

# ============================================
# 2. 工具创建
# ============================================

# Tavily 搜索工具实例
tavily_tool = TavilySearchResults(max_results=5)

# Python 执行工具（REPL）
repl = PythonREPL()

@tool
def python_repl(code: Annotated[str, "要执行的Python代码"]):
    """
    执行 Python 代码的工具。
    输出会显示给用户。如果想查看值，请使用 print()。
    """
    try:
        result = repl.run(code)
    except BaseException as e:
        return f"执行失败，错误: {repr(e)}"
    # 返回执行结果，并提示如果完成任务加 FINAL ANSWER
    return f"成功执行:\n```python\n{code}\n```\nStdout: {result}\n\n如果完成任务，请回复 FINAL ANSWER。"

# ============================================
# 3. 代理状态定义
# ============================================

class AgentState(TypedDict):
    """
    定义每个节点传递的状态。
    messages: 消息列表，可追加
    sender: 当前消息发送者标识，用于路由
    """
    messages: Annotated[Sequence[BaseMessage], operator.add]
    sender: str

# 辅助函数：代理处理节点逻辑
def agent_node(state, agent, name):
    """
    调用代理处理当前状态，并将输出转换为适合状态图的格式
    """
    result = agent.invoke(state)
    if not isinstance(result, ToolMessage):
        # 如果不是工具消息，转成 AIMessage
        result = AIMessage(**result.dict(exclude={"type", "name"}), name=name)
    return {
        "messages": [result],
        "sender": name,
    }

# ============================================
# 4. LLM 初始化
# ============================================

llm = ChatOpenAI(
    api_key=os.getenv("MODELSCOPE_API_KEY"),
    openai_api_base=os.getenv("MODELSCOPE_API_BASE"),
    model="deepseek-ai/DeepSeek-R1-0528",
    streaming=True,
    temperature=0
)

# ============================================
# 5. 代理节点创建
# ============================================

# 研究代理（搜索数据）
research_agent = create_agent(
    llm,
    [tavily_tool],
    system_message="你应该提供准确的数据供 chart_generator 使用。"
)
research_node = functools.partial(agent_node, agent=research_agent, name="Researcher")

# 图表生成代理（执行Python代码生成图表）
chart_agent = create_agent(
    llm,
    [python_repl],
    system_message="你展示的任何图表都将对用户可见。"
)
chart_node = functools.partial(agent_node, agent=chart_agent, name="chart_generator")

# 工具节点
tools = [tavily_tool, python_repl]
tool_node = ToolNode(tools)

# ============================================
# 6. 路由器函数
# ============================================

def router(state) -> Literal["call_tool", "__end__", "continue"]:
    """
    根据当前状态决定下一步执行哪个节点。
    - 如果上一个消息在调用工具 -> call_tool
    - 如果包含 FINAL ANSWER -> 结束
    - 否则继续执行下一个代理
    """
    messages = state["messages"]
    last_message = messages[-1]

    if last_message.tool_calls:
        return "call_tool"
    if "FINAL ANSWER" in last_message.content:
        return "__end__"
    return "continue"

# ============================================
# 7. 构建状态图
# ============================================

workflow = StateGraph(AgentState)

# 添加节点
workflow.add_node("Researcher", research_node)
workflow.add_node("chart_generator", chart_node)
workflow.add_node("call_tool", tool_node)

# 添加条件边（根据路由器返回值选择下一个节点）
workflow.add_conditional_edges(
    "Researcher",
    router,
    {
        "continue": "chart_generator",
        "call_tool": "call_tool",
        "__end__": END
    },
)

workflow.add_conditional_edges(
    "chart_generator",
    router,
    {
        "continue": "Researcher",
        "call_tool": "call_tool",
        "__end__": END
    },
)

# 工具节点根据 sender 返回到对应代理
workflow.add_conditional_edges(
    "call_tool",
    lambda x: x["sender"],
    {
        "Researcher": "Researcher",
        "chart_generator": "chart_generator"
    },
)

# 起始边
workflow.add_edge(START, "Researcher")

# 编译状态图
graph = workflow.compile()

# ============================================
# 8. 执行图
# ============================================

events = graph.stream(
    {
        "messages": [
            HumanMessage(
                content="获取过去5年AI软件市场规模，然后绘制一条折线图。一旦你编写好代码，完成任务。"
            )
        ],
    },
    {"recursion_limit": 150}  # 防止无限循环
)

# 打印事件流
for s in events:
    print(s)
    print("----")
